#include <iostream>
#include <hip/hip_runtime.h>
#include <thread>
#include <vector>
#include <chrono>
#include <cstdlib>
#include <algorithm>
#include <random>

void allocate_and_use_memory(int device_id, size_t increment_size, int duration) {
    // 设置当前设备
    hipSetDevice(device_id);

    std::vector<void*> allocations;
    std::vector<float*> used_memory; // 用于存储实际使用的内存指针

    std::cout << "Process " << device_id << " starting memory allocation." << std::endl;

    auto start_time = std::chrono::steady_clock::now();
    while (std::chrono::steady_clock::now() - start_time < std::chrono::seconds(duration)) {
        void* d_ptr;
        hipMalloc(&d_ptr, increment_size);
        allocations.push_back(d_ptr);

        // 实际使用内存（例如，初始化为某个值）
        float* h_ptr = new float[increment_size / sizeof(float)];
        std::fill(h_ptr, h_ptr + (increment_size / sizeof(float)), static_cast<float>(device_id));
        hipMemcpy(d_ptr, h_ptr, increment_size, hipMemcpyHostToDevice);
        used_memory.push_back(h_ptr); // 保存用于释放的指针

        std::cout << "Process " << device_id << ": Allocated and used " << increment_size / (1024 * 1024) << " MB" << std::endl;
        std::this_thread::sleep_for(std::chrono::seconds(1)); // 每秒申请一次内存
    }

    // 打乱释放顺序
    std::cout << "Process " << device_id << " will start releasing memory." << std::endl;
    std::shuffle(allocations.begin(), allocations.end(), std::default_random_engine(std::chrono::system_clock::now().time_since_epoch().count()));

    for (size_t i = 0; i < allocations.size(); ++i) {
        hipFree(allocations[i]);
        std::cout << "Process " << device_id << ": Released GPU memory chunk " << (i + 1) << std::endl;
    }

    // 释放主机内存
    for (size_t i = 0; i < used_memory.size(); ++i) {
        delete[] used_memory[i];
        std::cout << "Process " << device_id << ": Released CPU memory chunk " << (i + 1) << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(500)); // 每0.5秒释放一次内存
    }

    allocations.clear();
    used_memory.clear();
    std::cout << "Process " << device_id << " has released all allocated memory." << std::endl;
    return;
}

void allocate_and_use_cpu_memory(size_t increment_size, int duration) {
    std::vector<void*> allocations;
    std::cout << "CPU process starting memory allocation." << std::endl;

    auto start_time = std::chrono::steady_clock::now();
    while (std::chrono::steady_clock::now() - start_time < std::chrono::seconds(duration)) {
        void* ptr = malloc(increment_size);
        allocations.push_back(ptr);

        // 实际使用内存（例如，初始化为某个值）
        std::fill(static_cast<char*>(ptr), static_cast<char*>(ptr) + increment_size, static_cast<char>(1));

        std::cout << "CPU process: Allocated and used " << increment_size / (1024 * 1024) << " MB" << std::endl;
        std::this_thread::sleep_for(std::chrono::seconds(1)); // 每秒申请一次内存
    }

    // 打乱释放顺序
    std::cout << "CPU process will start releasing memory." << std::endl;
    std::shuffle(allocations.begin(), allocations.end(), std::default_random_engine(std::chrono::system_clock::now().time_since_epoch().count()));

    for (size_t i = 0; i < allocations.size(); ++i) {
        free(allocations[i]);
        std::cout << "CPU process: Released memory chunk " << (i + 1) << std::endl;
        std::this_thread::sleep_for(std::chrono::milliseconds(500)); // 每0.5秒释放一次内存
    }

    allocations.clear();
    std::cout << "CPU process has released all allocated memory." << std::endl;
    return;
}

int main() {
    int num_devices;
    hipGetDeviceCount(&num_devices); // 动态获取可用的GPU数量
    std::cout << "Number of available GPUs: " << num_devices << std::endl;

    const int duration = 30; // 持续时间30秒

    std::vector<std::thread> threads;

    // 启动GPU内存申请线程
    for (int device_id = 0; device_id < num_devices; ++device_id) {
        size_t increment_size = (device_id + 1) * 5 * 1024 * 1024; // 每张卡递增5M, 10M, 15M, ...
        threads.emplace_back(allocate_and_use_memory, device_id, increment_size, duration);
    }

    // 启动CPU内存申请线程
    threads.emplace_back(allocate_and_use_cpu_memory, 10 * 1024 * 1024, duration); // CPU每次递增10MB

    for (auto& thread : threads) {
        thread.join();
    }

    return 0;
}
